import os
import re
import sys
import logging
import tempfile
import subprocess
import time
import random
import json
import asyncio
import numpy as np
import torch
import torchaudio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, BackgroundTasks, HTTPException
from enum import Enum
from typing import List, Dict, Any, Optional
from queue import Queue
from model import SenseVoiceSmall
from funasr.utils.postprocess_utils import rich_transcription_postprocess

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("sensevoice-stream")

class Language(str, Enum):
    auto = "auto"
    zh = "zh"
    en = "en"
    yue = "yue"
    ja = "ja"
    ko = "ko"
    nospeech = "nospeech"

# WebSocket连接管理
class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.session_data: Dict[str, ASRSession] = {}
    
    async def connect(self, websocket: WebSocket, session_id: str) -> None:
        await websocket.accept()
        self.active_connections[session_id] = websocket
        self.session_data[session_id] = ASRSession(session_id)
        logger.info(f"WebSocket连接建立: {session_id}")
    
    def disconnect(self, session_id: str) -> None:
        if session_id in self.active_connections:
            del self.active_connections[session_id]
        if session_id in self.session_data:
            self.session_data[session_id].cleanup()
            del self.session_data[session_id]
        logger.info(f"WebSocket连接关闭: {session_id}")
    
    async def send_text(self, session_id: str, message: str) -> None:
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_text(message)
            except Exception as e:
                logger.error(f"发送文本消息失败: {str(e)}")
                self.disconnect(session_id)
    
    async def send_json(self, session_id: str, data: Dict[str, Any]) -> None:
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_json(data)
            except Exception as e:
                logger.error(f"发送JSON消息失败: {str(e)}")
                self.disconnect(session_id)
    
    def get_session(self, session_id: str) -> Optional['ASRSession']:
        return self.session_data.get(session_id)

# 音频处理和ASR会话
class ASRSession:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.audio_chunks: List[bytes] = []
        self.total_audio_bytes = 0
        self.processed_chunks = 0
        self.audio_format = "audio/webm"  # 默认格式
        self.language = "zh"  # 默认语言
        self.is_streaming = True
        self.previous_state = None  # 用于存储模型状态
        self.text_buffer = ""  # 存储已识别文本
        self.previous_text = ""  # 上一次识别的文本，用于比较差异
        self.temp_dir = "./temp_audio/stream"
        self.temp_files = []
        self.processing = False  # 处理锁
        self.last_process_time = time.time()
        self.started_at = time.time()  # 会话开始时间
        self.last_activity_time = time.time()  # 最后活动时间，用于超时检测
        self.min_audio_size = int(os.getenv("MIN_AUDIO_SIZE_FOR_PROCESSING", "4000"))  # 最小处理音频大小
        self.retry_count = 0  # 处理重试计数
        self.max_retries = 3  # 最大重试次数
        
        # 创建临时目录
        try:
            os.makedirs(self.temp_dir, exist_ok=True)
        except Exception as e:
            logger.error(f"创建临时目录失败: {str(e)}")
            raise RuntimeError(f"无法创建临时目录: {str(e)}")
            
    def add_audio_chunk(self, chunk: bytes, is_final: bool = False) -> None:
        """添加音频数据块"""
        if len(chunk) > 0:
            self.audio_chunks.append(chunk)
            self.total_audio_bytes += len(chunk)
            self.last_process_time = time.time()
            self.last_activity_time = time.time()
            logger.debug(f"添加音频块，大小: {len(chunk)}字节，累计: {self.total_audio_bytes}字节，是否最终: {is_final}")
            
        if is_final:
            logger.info(f"接收到最终标记，切换为非流式模式，会话ID: {self.session_id}")
            self.is_streaming = False
    
    def set_metadata(self, metadata: Dict[str, Any]) -> None:
        """设置会话元数据"""
        # 更新最后活动时间
        self.last_activity_time = time.time()
        if "format" in metadata:
            self.audio_format = metadata["format"]
        if "language" in metadata and metadata["language"] in [lang.value for lang in Language]:
            self.language = metadata["language"]
        else:
            logger.warning(f"不支持的语言: {metadata.get('language', 'None')}, 使用默认语言: {self.language}")
            
        if "streaming" in metadata:
            self.is_streaming = metadata["streaming"]
            # 当模式切换时，重置文本缓冲区
            if not self.is_streaming:
                self.text_buffer = ""
                self.previous_text = ""
        
        logger.info(f"设置元数据: 格式={self.audio_format}, 语言={self.language}, 流式模式={self.is_streaming}")
    
    async def process_audio(self, manager: ConnectionManager, model, model_kwargs) -> str:
        """处理累积的音频数据"""
        if self.processing or not self.audio_chunks:
            return ""
        # 检查累积的音频数据大小是否达到处理阈值
        accumulated_size = sum(len(chunk) for chunk in self.audio_chunks[self.processed_chunks:])
        if accumulated_size < self.min_audio_size and self.is_streaming:  # 流式模式下需要累积足够数据
            logger.debug(f"累积的音频数据未达到处理阈值，当前大小: {accumulated_size}字节，阈值: {self.min_audio_size}字节")
            return ""
       
        self.processing = True
        logger.info(f"开始处理音频数据，会话ID: {self.session_id}, 累积大小: {accumulated_size}字节, 流式模式: {self.is_streaming}")
        
        try:
            # 准备进行处理的音频块
            start_idx = self.processed_chunks
            end_idx = len(self.audio_chunks)
            chunks_to_process = self.audio_chunks[start_idx:end_idx]
            
            if not chunks_to_process:
                self.processing = False
                return ""
            
            # 合并音频块
            combined_audio = b''.join(chunks_to_process)
            
            # 获取临时文件路径
            timestamp = int(time.time() * 1000)
            random_suffix = random.randint(1000, 9999)
            input_path = f"{self.temp_dir}/input_{self.session_id}_{timestamp}_{random_suffix}.webm"
            output_path = f"{self.temp_dir}/output_{self.session_id}_{timestamp}_{random_suffix}.wav"
            
            self.temp_files.extend([input_path, output_path])
            
            # 保存二进制数据到临时文件
            with open(input_path, "wb") as f:
                f.write(combined_audio)
            
            logger.info(f"已保存音频数据到临时文件: {input_path}, 大小: {len(combined_audio)}字节")
            
            # 使用ffmpeg转换格式
            try:
                # 检查文件是否存在并且非空
                if not os.path.exists(input_path) or os.path.getsize(input_path) == 0:
                    logger.error(f"音频文件不存在或为空: {input_path}")
                    self.retry_count += 1
                    if self.retry_count <= self.max_retries:
                        logger.warning(f"重试处理音频 ({self.retry_count}/{self.max_retries})")
                        self.processing = False
                        return ""
                    else:
                        logger.error(f"达到最大重试次数，放弃处理")
                        await manager.send_json(self.session_id, {
                            "type": "error",
                            "message": "音频处理失败，达到最大重试次数",
                            "code": "MAX_RETRIES_REACHED",
                            "session_id": self.session_id
                        })
                        self.processing = False
                        return ""
                
                # 使用非常简单的ffmpeg命令，避免复杂处理
                ffmpeg_cmd = [
                    'ffmpeg',
                    '-y',  # 覆盖输出文件
                    '-i', input_path,  # 输入文件
                    '-acodec', 'pcm_s16le',  # 16位PCM编码
                    '-ar', '16000',  # 16kHz采样率
                    '-ac', '1',  # 单声道
                    output_path  # 输出文件
                ]
                
                logger.info(f"执行ffmpeg命令: {' '.join(ffmpeg_cmd)}")
                
                # 执行ffmpeg命令
                process = subprocess.run(
                    ffmpeg_cmd,
                    check=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    timeout=30  # 30秒超时
                )
                
                stderr = process.stderr.decode('utf-8', errors='ignore')
                if stderr:
                    logger.debug(f"ffmpeg输出: {stderr}")
                
                logger.info(f"音频转换成功: {output_path}, 输出文件大小: {os.path.getsize(output_path) if os.path.exists(output_path) else 0}字节")
                
                # 检查输出文件是否生成
                if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
                    logger.error(f"ffmpeg转换后文件不存在或为空: {output_path}")
                    
                    # 尝试使用备用命令
                    backup_cmd = [
                        'ffmpeg',
                        '-y',
                        '-f', 'webm',  # 强制指定输入格式
                        '-i', input_path,
                        '-f', 'wav',   # 强制指定输出格式
                        output_path
                    ]
                    
                    logger.info(f"尝试备用ffmpeg命令: {' '.join(backup_cmd)}")
                    
                    try:
                        subprocess.run(
                            backup_cmd,
                            check=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            timeout=30
                        )
                        logger.info(f"备用命令转换成功: {output_path}")
                    except Exception as e:
                        logger.error(f"备用命令也失败: {str(e)}")
                        
                        # 即使转换失败，也返回空结果而不是报错，确保通信流程完整
                        if not self.is_streaming:
                            logger.warning(f"非流式模式下转换失败，发送空结果, 会话ID: {self.session_id}")
                            await manager.send_json(self.session_id, {
                                "type": "recognition",
                                "text": "",
                                "isFinal": True,
                                "timestamp": int(time.time() * 1000),
                                "session_id": self.session_id,
                                "streaming": False
                            })
                            self.processing = False
                            return ""  # 返回空结果但不抛出异常
                        else:
                            self.processing = False
                            return ""
                
                # 加载音频数据
                try:
                    data_tensor, sample_rate = torchaudio.load(output_path)
                    
                    # 单声道处理
                    if data_tensor.dim() > 1 and data_tensor.size(0) > 1:
                        data_tensor = data_tensor.mean(0)
                    elif data_tensor.dim() > 1:
                        data_tensor = data_tensor.squeeze(0)
                    
                    # 检查采样率
                    if sample_rate != 16000:
                        logger.warning(f"采样率 {sample_rate}Hz 与期望的16000Hz不符，执行重采样")
                        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
                        data_tensor = resampler(data_tensor)
                        sample_rate = 16000
                    
                    logger.info(f"音频加载成功，形状: {data_tensor.shape}, 采样率: {sample_rate}Hz")
                    
                except Exception as e:
                    logger.error(f"加载音频文件失败: {str(e)}")
                    # 即使加载失败，对于非流式模式也返回空结果
                    if not self.is_streaming:
                        await manager.send_json(self.session_id, {
                            "type": "recognition",
                            "text": "",
                            "isFinal": True,
                            "timestamp": int(time.time() * 1000),
                            "session_id": self.session_id,
                            "streaming": False
                        })
                    self.processing = False
                    return ""
                
                # 准备执行推理
                keys = ["audio_stream"]
                audios = [data_tensor]
                
                # 记录时间和设置推理参数
                inference_start_time = time.time()
                
                # 设置推理参数
                inference_params = {
                    "data_in": audios,
                    "language": self.language,
                    "use_itn": True,  # 对所有模式启用文字规范化
                    "ban_emo_unk": False,
                    "key": keys,
                    "fs": sample_rate,
                    **model_kwargs
                }
                
                logger.info(f"开始执行推理, 会话ID: {self.session_id}, 流式模式: {self.is_streaming}")
                
                # 执行推理
                try:
                    result = model.inference(**inference_params)
                    inference_time = time.time() - inference_start_time
                    logger.info(f"推理完成，耗时: {inference_time:.2f}秒, 会话ID: {self.session_id}")
                    
                except Exception as e:
                    logger.error(f"推理过程出错: {str(e)}")
                    # 对于非流式模式，确保返回响应
                    if not self.is_streaming:
                        await manager.send_json(self.session_id, {
                            "type": "recognition",
                            "text": "",
                            "isFinal": True,
                            "timestamp": int(time.time() * 1000),
                            "session_id": self.session_id,
                            "streaming": False
                        })
                    self.processing = False
                    return ""
                
                # 更新已处理索引
                self.processed_chunks = end_idx
                
                # 处理推理结果
                processed_result = self._process_result(result)
                logger.info(f"处理后的结果: {processed_result}, 会话ID: {self.session_id}")
                
                # 判断是否有识别结果
                if processed_result and "result" in processed_result and processed_result["result"]:
                    recognized_text = processed_result["result"][0].get("text", "")
                    
                    # 保存之前的文本
                    self.previous_text = self.text_buffer
                    
                    # 计算新增文本
                    new_text = self._get_new_text(recognized_text)
                    
                    # 更新文本缓冲区
                    self.text_buffer = recognized_text
                    
                    # 准备发送识别结果
                    is_final = not self.is_streaming
                    response = {
                        "type": "recognition",
                        "text": self.text_buffer,  # 完整文本
                        "isFinal": is_final,
                        "timestamp": int(time.time() * 1000),
                        "session_id": self.session_id,
                        "streaming": self.is_streaming
                    }
                    
                    # 添加增量文本（如果有）
                    if new_text:
                        response["new_text"] = new_text
                        logger.debug(f"增量文本: [{new_text}], 会话ID: {self.session_id}")
                    
                    logger.info(f"发送识别结果: {response}, 会话ID: {self.session_id}")
                    await manager.send_json(self.session_id, response)
                    
                    return new_text
                else:
                    # 未识别到文本但需要返回结果
                    if not self.is_streaming:
                        logger.warning(f"未识别到文本，发送空结果，会话ID: {self.session_id}")
                        await manager.send_json(self.session_id, {
                            "type": "recognition",
                            "text": "",
                            "isFinal": True,
                            "timestamp": int(time.time() * 1000),
                            "session_id": self.session_id,
                            "streaming": False
                        })
                    
                    return ""
                    
            except subprocess.CalledProcessError as e:
                stderr = e.stderr.decode('utf-8', errors='ignore') if e.stderr else ""
                logger.error(f"ffmpeg命令执行失败: {e.returncode}, 输出: {stderr}")
                
                # 对于非流式模式，确保返回响应
                if not self.is_streaming:
                    await manager.send_json(self.session_id, {
                        "type": "recognition",
                        "text": "",
                        "isFinal": True,
                        "timestamp": int(time.time() * 1000),
                        "session_id": self.session_id,
                        "streaming": False
                    })
                
                self.processing = False
                return ""
                
            except Exception as e:
                logger.error(f"音频处理过程中发生错误: {str(e)}")
                
                # 对于非流式模式，确保返回响应
                if not self.is_streaming:
                    await manager.send_json(self.session_id, {
                        "type": "recognition",
                        "text": "",
                        "isFinal": True,
                        "timestamp": int(time.time() * 1000),
                        "session_id": self.session_id,
                        "streaming": False
                    })
                
                self.processing = False
                return ""
                
        except Exception as e:
            logger.error(f"处理音频数据时发生未知错误: {str(e)}")
            self.retry_count += 1
            
            # 对于非流式模式，确保返回响应
            if not self.is_streaming:
                await manager.send_json(self.session_id, {
                    "type": "recognition",
                    "text": "",
                    "isFinal": True,
                    "timestamp": int(time.time() * 1000),
                    "session_id": self.session_id,
                    "streaming": False
                })
            
            self.processing = False
            return ""
        
        finally:
            # 清理临时文件并重置处理状态
            self.processing = False
            self._cleanup_temp_files()
            logger.info(f"完成音频处理，会话ID: {self.session_id}")
        
        return ""
    
    def _get_new_text(self, current_text: str) -> str:
        """计算新增的文本部分"""
        if not self.previous_text:
            return current_text
        
        # 如果当前文本是先前文本的扩展
        if current_text.startswith(self.previous_text):
            return current_text[len(self.previous_text):]
        
        # 尝试找到部分重叠
        for i in range(min(len(self.previous_text), len(current_text))):
            overlap_len = len(self.previous_text) - i
            if current_text.startswith(self.previous_text[-overlap_len:]):
                return current_text[overlap_len:]
        
        # 如果文本完全不同，返回整个新文本
        return current_text
    
    def _process_result(self, res):
        """处理识别结果，添加原始文本和干净文本"""
        if len(res) == 0:
            return {"result": []}
            
        regex = r"<\|.*\|>"
        for item in res[0]:
            item["raw_text"] = item["text"]
            item["clean_text"] = re.sub(regex, "", item["text"], 0, re.MULTILINE)
            item["text"] = rich_transcription_postprocess(item["text"])
        
        return {"result": res[0]}
    
    def _cleanup_temp_files(self) -> None:
        """清理临时文件"""
        for file_path in self.temp_files:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
            except Exception as e:
                logger.warning(f"清理临时文件失败: {file_path}, 错误: {str(e)}")
        
        # 清空已处理的文件列表
        self.temp_files = []
    
    def cleanup(self) -> None:
        """清理会话资源"""
        self._cleanup_temp_files()
        # 清空音频缓冲区
        self.audio_chunks = []
        self.total_audio_bytes = 0
        self.text_buffer = ""
        logger.info(f"会话清理完成，ID: {self.session_id}, 会话持续时间: {time.time() - self.started_at:.2f}秒")

# 加载模型
def load_model():
    """加载语音识别模型"""
    # 设置模型路径
    model_dir = "./models/SenseVoiceSmall"
    
    try:
        # 尝试从本地目录加载模型
        if os.path.exists(model_dir):
            logger.info(f"从本地目录加载模型: {model_dir}")
            model, model_kwargs = SenseVoiceSmall.from_pretrained(
                model=model_dir, 
                device=os.getenv("SENSEVOICE_DEVICE", "cuda:0")
            )
        else:
            # 从modelscope加载模型
            logger.info("本地模型目录不存在，尝试从modelscope加载模型")
            model_fallback = "iic/SenseVoiceSmall"
            model, model_kwargs = SenseVoiceSmall.from_pretrained(
                model=model_fallback, 
                device=os.getenv("SENSEVOICE_DEVICE", "cuda:0")
            )
        
        model.eval()
        logger.info("模型加载成功")
        return model, model_kwargs
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        return None, None

# 异步处理音频
async def process_audio_periodically(session_id: str, manager: ConnectionManager, model, model_kwargs):
    """定期处理会话中的音频数据"""
    session = manager.get_session(session_id)
    if not session:
        return
    
    try:
        # 处理累积的音频数据
        result = await session.process_audio(manager, model, model_kwargs)
        
        # 如果是非流式模式而且没有结果，确保发送一个空结果
        if not session.is_streaming and not result and not session.text_buffer:
            logger.warning(f"非流式模式处理后没有结果，发送空结果，会话ID: {session_id}")
            await manager.send_json(session_id, {
                "type": "recognition",
                "text": "",
                "isFinal": True,
                "timestamp": int(time.time() * 1000),
                "session_id": session_id,
                "streaming": False
            })
    except Exception as e:
        logger.error(f"定期处理音频失败: {str(e)}")
        # 尝试向客户端发送错误信息
        try:
            await manager.send_json(session_id, {
                "type": "error",
                "message": f"处理音频失败: {str(e)}",
                "code": "PROCESSING_ERROR",
                "session_id": session_id
            })
        except Exception:
            pass

async def cleanup_inactive_sessions():
    """定期清理不活跃的会话"""
    timeout_seconds = int(os.getenv("SESSION_TIMEOUT_SECONDS", "300"))
    while True:
        try:
            current_time = time.time()
            inactive_sessions = []
            
            for session_id, session in manager.session_data.items():
                # 计算会话不活跃时间
                inactive_time = current_time - session.last_activity_time
                if inactive_time > timeout_seconds:
                    inactive_sessions.append(session_id)
                    logger.info(f"检测到不活跃会话: {session_id}, 不活跃时间: {inactive_time:.2f}秒")
            
            # 清理不活跃的会话
            for session_id in inactive_sessions:
                logger.info(f"清理不活跃会话: {session_id}")
                manager.disconnect(session_id)
            
            # 每分钟检查一次
            await asyncio.sleep(60)
        except Exception as e:
            logger.error(f"清理会话时出错: {str(e)}")
            await asyncio.sleep(60)  # 出错后仍继续

# 创建FastAPI应用
app = FastAPI(title="SenseVoice Stream API", description="流式语音识别API服务")
manager = ConnectionManager()

# 全局模型变量
model = None
model_kwargs = None

@app.on_event("startup")
async def startup_event():
    """应用启动时加载模型并启动会话清理任务"""
    global model, model_kwargs
    model, model_kwargs = load_model()
    if not model:
        logger.error("模型加载失败，应用将无法正常工作")
        sys.exit(1)
    
    # 设置默认环境变量
    if not os.getenv("MIN_AUDIO_SIZE_FOR_PROCESSING"):
        os.environ["MIN_AUDIO_SIZE_FOR_PROCESSING"] = "4000"  # 最小处理音频大小 (字节)
    if not os.getenv("SESSION_TIMEOUT_SECONDS"):
        os.environ["SESSION_TIMEOUT_SECONDS"] = "300"  # 会话超时时间 (秒)
    if not os.getenv("PROCESSING_INTERVAL_MS"):
        os.environ["PROCESSING_INTERVAL_MS"] = "500"  # 处理间隔 (毫秒)
    
    logger.info(f"最小处理音频大小: {os.getenv('MIN_AUDIO_SIZE_FOR_PROCESSING')}字节")
    logger.info(f"会话超时时间: {os.getenv('SESSION_TIMEOUT_SECONDS')}秒")
    logger.info(f"处理间隔: {os.getenv('PROCESSING_INTERVAL_MS')}毫秒")
    
    # 启动会话清理任务
    asyncio.create_task(cleanup_inactive_sessions())

@app.get("/")
async def root():
    """API根路径"""
    return {
        "name": "SenseVoice Stream API",
        "version": "1.0.0",
        "status": "running",
        "endpoints": [
            "/api/v1/asr/stream"
        ]
    }

@app.get("/health")
async def health():
    """健康检查端点"""
    if model is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    return {"status": "healthy", "active_sessions": len(manager.session_data)}

@app.websocket("/api/v1/asr/stream")
async def websocket_stream(websocket: WebSocket):
    """WebSocket流式ASR端点"""
    session_id = f"session_{int(time.time())}_{random.randint(1000, 9999)}"
    
    await manager.connect(websocket, session_id)
    
    # 发送会话ID给客户端
    await manager.send_json(session_id, {
        "type": "connection",
        "sessionId": session_id,
        "message": "连接成功"
    })
    
    # 创建定期处理任务
    processing_task = None
    
    # 周期性处理间隔 (毫秒)
    processing_interval_ms = int(os.getenv("PROCESSING_INTERVAL_MS", "500"))
    
    try:
        while True:
            # 接收消息
            message = await websocket.receive()
            
            # 处理文本消息
            if "text" in message:
                try:
                    data = json.loads(message["text"])
                    msg_type = data.get("type")
                    
                    if msg_type == "audio_metadata":
                        # 处理元数据
                        session = manager.get_session(session_id)
                        if session:
                            metadata = {
                                "format": data.get("format", "audio/webm"),
                                "language": data.get("lang", "zh"),
                                "streaming": data.get("streaming", True)
                            }
                            session.set_metadata(metadata)
                            logger.info(f"接收到元数据: {metadata}, 会话ID: {session_id}")
                            
                            # 确认收到元数据
                            await manager.send_json(session_id, {
                                "type": "metadata_ack",
                                "status": "received",
                                "session_id": session_id
                            })
                    elif msg_type == "end":
                        # 处理结束标记
                        session = manager.get_session(session_id)
                        if session:
                            session.add_audio_chunk(b'', is_final=True)
                            logger.info(f"接收到结束标记，会话ID: {session_id}")
                            
                            # 特别处理非流式模式下的结束标记
                            if not session.is_streaming:
                                logger.info(f"非流式模式下收到结束标记，立即处理最终结果，会话ID: {session_id}")
                                try:
                                    await process_audio_periodically(session_id, manager, model, model_kwargs)
                                    # 确保有结果返回，即使是空结果
                                    if not session.text_buffer:
                                        logger.warning(f"非流式模式下没有识别结果，发送空结果，会话ID: {session_id}")
                                        await manager.send_json(session_id, {
                                            "type": "recognition",
                                            "text": "",
                                            "isFinal": True,
                                            "timestamp": int(time.time() * 1000),
                                            "session_id": session_id,
                                            "streaming": False
                                        })
                                except Exception as e:
                                    logger.error(f"处理最终结果失败: {str(e)}")
                            
                            # 确认收到结束标记
                            await manager.send_json(session_id, {
                                "type": "end_ack",
                                "status": "received",
                                "session_id": session_id
                            })
                    
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析错误: {str(e)}")
                    await manager.send_json(session_id, {
                        "type": "error",
                        "message": f"无效的JSON格式: {str(e)}",
                        "code": "INVALID_JSON",
                        "session_id": session_id
                    })
                except Exception as e:
                    logger.error(f"处理文本消息失败: {str(e)}")
                    await manager.send_json(session_id, {
                        "type": "error",
                        "message": f"处理文本消息失败: {str(e)}",
                        "code": "TEXT_PROCESSING_ERROR",
                        "session_id": session_id
                    })
            
            # 处理二进制消息
            elif "bytes" in message:
                try:
                    binary_data = message["bytes"]
                    
                    # 获取会话
                    session = manager.get_session(session_id)
                    if not session:
                        continue
                    
                    # 添加音频块
                    is_final = False 
                    
                    # 对于非流式处理，每个二进制消息后都立即处理
                    if not session.is_streaming:
                        is_final = True
                        logger.info(f"非流式模式收到音频数据，标记为最终数据，会话ID: {session_id}, 大小: {len(binary_data)}字节")
                    
                    session.add_audio_chunk(binary_data, is_final)
                    
                    # 确认接收到数据
                    if len(binary_data) > 1000:  # 只对大于1KB的数据块发送确认
                        await manager.send_json(session_id, {
                            "type": "data_ack",
                            "size": len(binary_data),
                            "total": session.total_audio_bytes,
                            "session_id": session_id
                        })
                    
                    # 按需启动处理任务
                    if processing_task is None or processing_task.done():
                        processing_task = asyncio.create_task(
                            process_audio_periodically(session_id, manager, model, model_kwargs)
                        )
                        
                    # 非流式模式下，接收到数据后立即处理
                    if not session.is_streaming:
                        try:
                            # 等待处理完成
                            await processing_task
                            
                            # 如果处理后仍无结果，确保发送空结果，避免前端卡住
                            if not session.text_buffer:
                                logger.warning(f"非流式处理无结果，发送空结果，会话ID: {session_id}")
                                await manager.send_json(session_id, {
                                    "type": "recognition",
                                    "text": "",
                                    "isFinal": True,
                                    "timestamp": int(time.time() * 1000),
                                    "session_id": session_id,
                                    "streaming": False
                                })
                        except Exception as e:
                            logger.error(f"立即处理音频失败: {str(e)}")
                            # 确保即使处理失败也返回结果给前端
                            await manager.send_json(session_id, {
                                "type": "error",
                                "message": f"处理音频失败: {str(e)}",
                                "code": "AUDIO_PROCESSING_ERROR",
                                "session_id": session_id
                            })
                    
                except Exception as e:
                    logger.error(f"处理二进制消息失败: {str(e)}")
                    await manager.send_json(session_id, {
                        "type": "error",
                        "message": f"处理二进制消息失败: {str(e)}",
                        "code": "BINARY_PROCESSING_ERROR",
                        "session_id": session_id
                    })
    
    except WebSocketDisconnect:
        logger.info(f"WebSocket连接断开: {session_id}")
        
        # 发送最后的处理结果
        session = manager.get_session(session_id)
        if session and session.text_buffer:
            try:
                # 尝试进行最后一次处理
                if processing_task:
                    try:
                        await processing_task
                    except asyncio.CancelledError:
                        pass
                
                # 发送最终结果
                await manager.send_json(session_id, {
                    "type": "recognition",
                    "text": session.text_buffer,
                    "isFinal": True,
                    "timestamp": int(time.time() * 1000),
                    "session_id": session_id,
                    "streaming": False,
                    "message": "连接断开，处理完成"
                })
            except Exception as e:
                logger.error(f"发送最终结果失败: {str(e)}")
                
    except Exception as e:
        logger.error(f"WebSocket处理错误: {str(e)}")
    finally:
        # 清理资源
        manager.disconnect(session_id)
        if processing_task and not processing_task.done():
            try:
                processing_task.cancel()
                await asyncio.sleep(0.1)  # 给任务一点时间来取消
            except Exception as e:
                logger.error(f"取消处理任务失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    
    # 启动服务
    port = int(os.getenv("PORT", 9001))
    logger.info(f"启动SenseVoice流式语音识别API服务，访问地址: ws://0.0.0.0:{port}/api/v1/asr/stream")
    uvicorn.run(app, host="0.0.0.0", port=port) 